# include "utils.h"

namespace mslite {
    
double getAvg(std::vector<double> vals, int removes) {
    sort(vals.begin(), vals.end());
    double avg = 0;
    int cnt = 0;
    for (int i=removes; i<vals.size()-removes; i++) {
        avg += vals[i];
        cnt++;
    }
    return (avg / cnt);
}

int readSrcCode(const std::string filename, std::string &srcCode) {
    std::stringstream ss;
    std::ifstream clFile(filename, std::ios::in);
    ss << clFile.rdbuf();
    clFile.close();
    srcCode = std::move(ss.str());
    return 0;
}

int getTotalSize(std::vector<int> shape) {
    if (shape.size() == 0)
        return 0;
    int sizeTotal = 1;
    for (auto dim: shape) {
        sizeTotal *= dim;
    }
    return sizeTotal;
}

int genData(std::vector<float>& data, std::vector<int> shape) {
    if (shape.size() == 0) {
        data.clear();
        return 0;
    }
    int seed = time(0);
    // std::cout << seed << std::endl;
    std::default_random_engine generator(seed);
    std::uniform_real_distribution<float> urd(-1.0, 1.0);
    int size = 1;
    for (int i=0; i<shape.size(); i++) {
        size *= shape[i];
    }
    data.resize(size);
    for (int i=0; i<size; i++) {
        data[i] = urd(generator);
        // data[i] = 1;
    }
    return 0;
}

// lhs - m*n,  rhs - n*k
std::vector<float>&& matmul(std::vector<float> &lhs, std::vector<float> &rhs, int m, int n, int k) {
    std::vector<float> res(m*k);
    for (int y = 0; y < m; ++y) {
        for (int x = 0; x < k; ++x) {
            float sum = 0.0f;
            for (int i = 0; i < n; ++i) {
                sum += lhs[y * n + i] * rhs[i * k + x];
            }
            res[y * k + x] = sum;
        }
    }
    return std::move(res);
}

std::vector<float>&& RearrangeWeightForWino(std::vector<float> &weight, std::vector<int> shape) {
    std::vector<float> G, GT;
    getGAndGt(G, GT);

    int oc = shape[0], ic = shape[3];
    std::vector<float> res(oc*6*6*ic);
    for (int o=0; o<oc; o++) {
        for (int i=0; i<ic; i++) {
            std::vector<float> invals(9);
            for (int y=0; y<3; y++) 
                for (int x=0; x<3; x++) 
                    invals[y*3+x] = weight[((o*3+y)*3+x)*ic+i];

            std::vector<float> temp_vals = matmul(G, invals, 6, 3, 3);
            std::vector<float> out_vals = matmul(temp_vals, GT, 6, 3, 6);
            for (int y=0; y<6; y++) 
                for (int x=0; x<6; x++) 
                    res[((o*6+y)*6+x)*ic+i] = out_vals[y*6+x];
        }
    }
    return std::move(res);
}

// eg: ogroup=1, 128*6*6*128 --> 32*6*6*32*4*4
//     ogroup=2, 128*6*6*128 --> 16*6*6*32*2*4*4
//     ogroup=4, 128*6*6*128 --> 8*6*6*32*4*4*4
std::vector<float> ReorderOHWIToI4O4(const std::vector<float> &weight, std::vector<int> shape, int ogroup) {
    std::vector<float> res(weight.size());
    int oc = shape[0], ic = shape[3], kh = shape[1], kw = shape[2];
    int oc_slice_grp = oc / (4*ogroup), ic_slice = ic / 4, cnt = 0;
    for (int osg=0; osg<oc_slice_grp; osg++) {
        for (int y=0; y<kh; y++) {
            for (int x=0; x<kw; x++) {
                for (int is=0; is<ic_slice; is++) {
                    for (int og=0; og<ogroup; og++) {
                        for (int j=0; j<4; j++) {
                            for (int i=0; i<4; i++) {
                                int cur_ic = is*4 + j;
                                int cur_oc = (osg * ogroup + og) * 4 + i;
                                res[cnt++] = weight[((cur_oc*6+y)*6+x)*ic+cur_ic];
                            }
                        }
                    }
                }
            }
        }
    }
    // return std::move(res);
    return res;
}

void getGAndGt(std::vector<float> G, std::vector<float> GT) {
    GT.resize(6*3);
    G.resize(6*3);
    GT = {1.000000000000,  1.000000000000,  1.000000000000, 1.000000000000,  1.000000000000,  0.000000000000,
          0.000000000000,  0.707106769085, -0.707106769085, 1.414213538170, -1.414213538170,  0.000000000000,
          0.000000000000,  0.499999970198,  0.499999970198, 1.999999880791,  1.999999880791,  1.000000000000};
    for (int y = 0; y < 3; ++y) {
        for (int x = 0; x < 6; ++x) {
            G[x * 3 + y] = GT[y * 6 + x];
        }
    }
}

}
